#ifndef MEMORY_POOL_H__
#define MEMORY_POOL_H__

#include "define.h"
#include <assert.h>

namespace engine
{
#define MEMORY_POOL_MAX_FREE_PAGES 4

	template <class T> class MemoryPool
	{
	public:
		struct Page;
		struct MemoryWithPage
		{
			T userMemory;
			Page *parentPage;
		};

		struct Page
		{
			MemoryWithPage** availableStack;
			int availableStackSize;
			MemoryWithPage* block;
			Page *next, *prev;
		};

		MemoryPool();

		~MemoryPool();
		
		void setPageSize(int32 size);
		
		T *alloc();
		
		void release(T* ptr);
		
		void clear();

	protected:
		int blocksPerPage() const;
		
		bool initPage(Page *page, Page *prev);

		Page* availablePages;
		Page* unavailablePages;
		int32 availablePagesSize;
		int32 unavailablePagesSize;
		int32 memoryPoolPageSize;
	};

	template<class T>
	MemoryPool<T>::MemoryPool()
	{
		availablePagesSize = 0;
		unavailablePagesSize = 0;
		memoryPoolPageSize = 16384;
	}

	template<class T>
	MemoryPool<T>::~MemoryPool()
	{
		clear();
	}

	template<class T>
	void MemoryPool<T>::setPageSize(int size)
	{
		memoryPoolPageSize = size;
	}

	template<class T>
	T* MemoryPool<T>::alloc()
	{
		if (availablePagesSize > 0)
		{
			Page* curPage = availablePages;
			T* retVal = (T*)curPage->availableStack[--(curPage->availableStackSize)];

			if (curPage->availableStackSize == 0)
			{
				--availablePagesSize;
				availablePages = curPage->next;
				assert(availablePagesSize == 0 || availablePages->availableStackSize > 0);
				curPage->next->prev = curPage->prev;
				curPage->prev->next = curPage->next;

				if (unavailablePagesSize++ == 0)
				{
					unavailablePages = curPage;
					curPage->next = curPage;
					curPage->prev = curPage;
				}
				else
				{
					curPage->next = unavailablePages;
					curPage->prev = unavailablePages->prev;
					unavailablePages->prev->next = curPage;
					unavailablePages->prev = curPage;
				}
			}

			assert(availablePagesSize == 0 || availablePages->availableStackSize > 0);
			
			return retVal;
		}

		availablePages = (Page*)malloc(sizeof(Page));
		
		if (!availablePages)
			return NULL;
		
		availablePagesSize = 1;
		
		if (!initPage(availablePages, availablePages))
			return NULL;

		assert(availablePages->availableStackSize > 1);

		return (T*)availablePages->availableStack[--availablePages->availableStackSize];
	}

	template<class T>
	void MemoryPool<T>::release(T* ptr)
	{
		// Find the page this block is in and return it.
		Page *curPage;
		MemoryWithPage* memoryWithPage = (MemoryWithPage*)ptr;
		curPage = memoryWithPage->parentPage;

		if (curPage->availableStackSize == 0)
		{
			// The page is in the unavailable list so move it to the available list
			curPage->availableStack[curPage->availableStackSize++] = memoryWithPage;
			unavailablePagesSize--;

			// As this page is no longer totally empty, move it to the end of available pages
			curPage->next->prev = curPage->prev;
			curPage->prev->next = curPage->next;

			if (unavailablePagesSize > 0 && curPage == unavailablePages)
				unavailablePages = unavailablePages->next;

			if (availablePagesSize++ == 0)
			{
				availablePages = curPage;
				curPage->next = curPage;
				curPage->prev = curPage;
			}
			else
			{
				curPage->next = availablePages;
				curPage->prev = availablePages->prev;
				availablePages->prev->next = curPage;
				availablePages->prev = curPage;
			}
		}
		else
		{
			curPage->availableStack[curPage->availableStackSize++] = memoryWithPage;

			if (curPage->availableStackSize == blocksPerPage() && availablePagesSize >= MEMORY_POOL_MAX_FREE_PAGES)
			{
				// After a certain point, just deallocate empty pages rather than keep them around
				if (curPage == availablePages)
				{
					availablePages = curPage->next;
					assert(availablePages->availableStackSize > 0);
				}
				
				curPage->prev->next = curPage->next;
				curPage->next->prev = curPage->prev;
				availablePagesSize--;
				
				free(curPage->availableStack);
				free(curPage->block);
				free(curPage);
			}
		}
	}
	template<class T>
	void MemoryPool<T>::clear(void)
	{
		Page *cur, *freed;

		if (availablePagesSize>0)
		{
			cur = availablePages;
			
			while (true)
			{
				free(cur->availableStack);
				free(cur->block);
				freed = cur;
				cur = cur->next;
				if (cur == availablePages)
				{
					free(freed);
					break;
				}
				free(freed);
			}
		}

		if (unavailablePagesSize>0)
		{
			cur = unavailablePages;
			while (true)
			{
				free(cur->availableStack);
				free(cur->block);
				freed = cur;
				cur = cur->next;
				if (cur == unavailablePages)
				{
					free(freed);
					break;
				}
				free(freed);
			}
		}

		availablePagesSize = 0;
		unavailablePagesSize = 0;
	}

	template<class T>
	int MemoryPool<T>::blocksPerPage() const
	{
		return memoryPoolPageSize / sizeof(MemoryWithPage);
	}

	template<class T>
	bool MemoryPool<T>::initPage(Page *page, Page *prev)
	{
		int i = 0;
		const int bpp = blocksPerPage();
		page->block = (MemoryWithPage*)malloc(memoryPoolPageSize);
		if (!page->block)
			return false;

		page->availableStack = (MemoryWithPage**)malloc(sizeof(MemoryWithPage*) * bpp);
		if (!page->availableStack)
		{
			free(page->block);
			return false;
		}
		MemoryWithPage* curBlock = page->block;
		MemoryWithPage** curStack = page->availableStack;
		
		while (i < bpp)
		{
			curBlock->parentPage = page;
			curStack[i] = curBlock++;
			i++;
		}
		page->availableStackSize = bpp;
		page->next = availablePages;
		page->prev = prev;

		return true;
	}
}


#endif
